`timescale 1ns / 1ps

module aru_reduce_tb;

    // 时钟和复位
    logic clk;
    logic rst_n;

    // 参数定义
    parameter int P_ARU = 4;
    parameter int N0 = 8;
    parameter int DATA_WIDTH = P_ARU * N0;  // 32

    // 时钟生成
    initial begin
        clk = 0;
        forever #5 clk = ~clk;  // 10ns周期，100MHz
    end

    // 接口实例化
    aru_reduce_cfg_if u_cfg_if (
        clk,
        rst_n
    );
    aru_payload_if u_payload_if (
        clk,
        rst_n
    );  // 输入接口
    aru_reduce_pld_if u_output_if (
        clk,
        rst_n
    );  // 输出接口

    // DUT实例化 - 顶层aru_reduce模块
    aru_reduce u_dut (
        .clk                (clk),
        .rst_n              (rst_n),
        .u_aru_cfg_if       (u_cfg_if.in),
        .u_aru_payload_if   (u_payload_if.in),
        .u_aru_reduce_pld_if(u_output_if.out)
    );

    // 测试任务和函数

    // 复位任务
    task reset();
        rst_n              = 0;
        u_cfg_if.vld       = 0;
        u_cfg_if.reduce_m  = 0;
        u_cfg_if.reduce_n  = 0;
        u_cfg_if.reduce_op = 0;
        u_cfg_if.slice_m   = 0;
        u_cfg_if.slice_n   = 0;
        u_cfg_if.instr_idx = 0;
        u_payload_if.vld   = 0;
        u_payload_if.dat   = '0;
        u_payload_if.sdb   = '0;
        u_output_if.rdy    = 1;  // 输出端准备好接收
        repeat (10) @(posedge clk);
        rst_n = 1;
        repeat (5) @(posedge clk);
        $display("[%0t] Reset completed", $time);
    endtask

    // 配置任务
    task send_config(input logic reduce_m, input logic reduce_n, input logic [1:0] reduce_op, input int slice_m,
                     input int slice_n, input int instr_idx);
        while (!u_cfg_if.rdy) @(posedge clk);

        @(posedge clk);
        u_cfg_if.vld       = 1;
        u_cfg_if.reduce_m  = reduce_m;
        u_cfg_if.reduce_n  = reduce_n;
        u_cfg_if.reduce_op = reduce_op;
        u_cfg_if.slice_m   = slice_m;
        u_cfg_if.slice_n   = slice_n;
        u_cfg_if.instr_idx = instr_idx;

        // 等待握手成功
        while (!u_cfg_if.rdy) @(posedge clk);

        // 握手成功后，下一个时钟周期拉低valid
        @(posedge clk);
        u_cfg_if.vld = 0;

        $display(
            "[%0t] Config sent: reduce_m=%0d, reduce_n=%0d, reduce_op=%0d, slice_m=%0d, slice_n=%0d, instr_idx=%0d",
            $time, reduce_m, reduce_n, reduce_op, slice_m, slice_n, instr_idx);
    endtask

    // 发送多周期数据任务
    task send_data_multi_cycles(input int start_value);
        real             float_value;
        int              row_value;
        shortreal        temp_float;
        logic     [31:0] float_bits;
        int              idx;

        row_value = start_value;

        for (int cyc = 0; cyc < 5; cyc++) begin
            while (!u_payload_if.rdy) @(posedge clk);

            @(posedge clk);
            u_payload_if.vld = 1;

            // 最后一个周期：设置eom和eon
            if (cyc == 4) begin
                u_payload_if.sdb.eom   = 1;
                u_payload_if.sdb.eon   = 1;
                u_payload_if.sdb.vld_m = 4;
                u_payload_if.sdb.vld_n = 4;
            end else begin
                u_payload_if.sdb.eom   = 0;
                u_payload_if.sdb.eon   = 0;
                u_payload_if.sdb.vld_m = 4;
                u_payload_if.sdb.vld_n = 8;
            end

            // 按行填充数据：P_ARU行(4行)，每行N0个数据(8列)
            for (int row = 0; row < P_ARU; row++) begin  // 4行
                float_value = $itor(row_value);  // 当前行的浮点值
                temp_float  = float_value;
                float_bits  = $shortrealtobits(temp_float);

                for (int col = 0; col < N0; col++) begin  // 每行8个数据
                    idx                            = row * N0 + col;  // 计算在dat数组中的索引
                    // BF16格式：取FP32的高16位（符号位+8位指数+7位尾数）
                    u_payload_if.dat.dat[idx].sign = float_bits[31];
                    u_payload_if.dat.dat[idx].exp  = float_bits[30:23];
                    u_payload_if.dat.dat[idx].mant = float_bits[22:16];
                end

                row_value = row_value + 1;  // 下一行的值递增
            end

            $display("[%0t] Input Cycle %0d/5: Data sent (row values: %.1f - %.1f), eom=%0d, eon=%0d", $time, cyc + 1,
                     $itor(row_value - P_ARU), $itor(row_value - 1), u_payload_if.sdb.eom, u_payload_if.sdb.eon);
        end

        while (!u_payload_if.rdy) @(posedge clk);

        @(posedge clk);
        u_payload_if.vld       = 0;
        u_payload_if.sdb.eom   = 0;
        u_payload_if.sdb.eon   = 0;
        u_payload_if.sdb.vld_m = 0;
        u_payload_if.sdb.vld_n = 0;

        $display("[%0t] 5-cycle transmission completed (row value range: %.1f - %.1f)", $time, $itor(start_value),
                 $itor(row_value - 1));
    endtask

    // 等待输出任务
    task wait_output(input int num_cycles);
        repeat (num_cycles) @(posedge clk);
    endtask

    // 测试场景

    // 测试1: MAX操作
    task test_max_operation();
        $display("\n========== Test 1: MAX Operation ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36), .instr_idx(1));
        send_data_multi_cycles(.start_value(1));
        wait_output(50);
        $display("Test 1 completed\n");
    endtask

    // 测试2: MIN操作
    task test_min_operation();
        $display("\n========== Test 2: MIN Operation ==========");
        send_config(.reduce_m(0), .reduce_n(1), .reduce_op(2'b01), .slice_m(4), .slice_n(36), .instr_idx(2));
        send_data_multi_cycles(.start_value(1));
        wait_output(50);
        $display("Test 2 completed\n");
    endtask

    // 测试3: SUM操作
    task test_sum_operation();
        $display("\n========== Test 3: SUM Operation ==========");
        send_config(.reduce_m(0), .reduce_n(1), .reduce_op(2'b10), .slice_m(4), .slice_n(36), .instr_idx(3));
        send_data_multi_cycles(.start_value(1));
        wait_output(50);
        $display("Test 3 completed\n");
    endtask

    // 测试4: AVG操作
    task test_avg_operation();
        $display("\n========== Test 4: AVG Operation ==========");
        send_config(.reduce_m(1), .reduce_n(1), .reduce_op(2'b11), .slice_m(4), .slice_n(36), .instr_idx(4));
        send_data_multi_cycles(.start_value(1));
        wait_output(50);
        $display("Test 4 completed\n");
    endtask

    // 测试5: BYPASS模式
    task test_bypass_mode();
        $display("\n========== Test 5: BYPASS Mode ==========");
        send_config(.reduce_m(0), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36), .instr_idx(5));
        send_data_multi_cycles(.start_value(1));
        wait_output(50);
        $display("Test 5 completed\n");
    endtask

    // 测试6: 多次连续操作
    task test_continuous_operations();
        $display("\n========== Test 6: Multiple Continuous Operations ==========");
        send_config(.reduce_m(1), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36), .instr_idx(1));
        send_data_multi_cycles(.start_value(1));
        // wait_output(10);

        send_config(.reduce_m(0), .reduce_n(1), .reduce_op(2'b01), .slice_m(4), .slice_n(36), .instr_idx(2));
        send_data_multi_cycles(.start_value(1));
        // wait_output(10);

        send_config(.reduce_m(0), .reduce_n(1), .reduce_op(2'b10), .slice_m(4), .slice_n(36), .instr_idx(3));
        send_data_multi_cycles(.start_value(1));
        // wait_output(10);

        send_config(.reduce_m(1), .reduce_n(1), .reduce_op(2'b11), .slice_m(4), .slice_n(36), .instr_idx(4));
        send_data_multi_cycles(.start_value(1));
        // wait_output(10);

        send_config(.reduce_m(0), .reduce_n(0), .reduce_op(2'b00), .slice_m(4), .slice_n(36), .instr_idx(5));
        send_data_multi_cycles(.start_value(1));
        wait_output(10);
        $display("Test 6 completed\n");
    endtask

    // 主测试流程
    initial begin
        $display("========== Simulation Start ==========");
        $display("Testing: Complete aru_reduce System");
        $display("Pipeline: Transpose -> Stage1 -> Stage2 -> Stage3 -> Div");
        $display("Data Send Rule: 5 cycles per test, last cycle with eom/eon=1");
        $display("Data Value Rule: Starting from 1.0, auto-increment by 1.0\n");

        reset();

        test_max_operation();
        test_min_operation();
        test_sum_operation();
        test_avg_operation();
        test_bypass_mode();
        test_continuous_operations();

        repeat (50) @(posedge clk);

        $display("========== Simulation End ==========");
        $finish;
    end

    // 超时保护
    initial begin
        #200000;
        $display("ERROR: Simulation timeout!");
        $finish;
    end

    // 波形dump
    initial begin
        $fsdbDumpfile("./aru_reduce.fsdb");
        $fsdbDumpvars("+all");
    end

endmodule
